//
// Copyright (c) 2023-2024 Ivica Siladic, Bruno Iljazovic, Korina Simicevic
//
// Distributed under the Boost Software License, Version 1.0.
// (See accompanying file LICENSE or copy at http://www.boost.org/LICENSE_1_0.txt)
//

#ifndef ASYNC_MQTT5_SUBSCRIBE_OP_HPP
#define ASYNC_MQTT5_SUBSCRIBE_OP_HPP

#include <cstdint>
#include <memory>
#include <string>
#include <vector>

#include <boost/asio/associated_allocator.hpp>
#include <boost/asio/associated_cancellation_slot.hpp>
#include <boost/asio/associated_executor.hpp>
#include <boost/asio/cancellation_type.hpp>
#include <boost/asio/error.hpp>
#include <boost/asio/detached.hpp>
#include <boost/asio/prepend.hpp>

#include <async_mqtt5/error.hpp>
#include <async_mqtt5/reason_codes.hpp>
#include <async_mqtt5/types.hpp>

#include <async_mqtt5/detail/cancellable_handler.hpp>
#include <async_mqtt5/detail/control_packet.hpp>
#include <async_mqtt5/detail/internal_types.hpp>
#include <async_mqtt5/detail/topic_validation.hpp>

#include <async_mqtt5/impl/codecs/message_decoders.hpp>
#include <async_mqtt5/impl/codecs/message_encoders.hpp>

#include <async_mqtt5/impl/disconnect_op.hpp>

namespace async_mqtt5::detail {

namespace asio = boost::asio;

template <typename ClientService, typename Handler>
class subscribe_op {
	using client_service = ClientService;

	struct on_subscribe {};
	struct on_suback {};

	std::shared_ptr<client_service> _svc_ptr;

	using handler_type = cancellable_handler<
		Handler,
		typename client_service::executor_type
	>;
	handler_type _handler;

	size_t _num_topics { 0 };

public:
	subscribe_op(
		std::shared_ptr<client_service> svc_ptr,
		Handler&& handler
	) :
		_svc_ptr(std::move(svc_ptr)),
		_handler(std::move(handler), _svc_ptr->get_executor())
	{
		auto slot = asio::get_associated_cancellation_slot(_handler);
		if (slot.is_connected())
			slot.assign([&svc = *_svc_ptr](asio::cancellation_type_t) {
				svc.cancel();
			});
	}

	subscribe_op(subscribe_op&&) = default;
	subscribe_op(const subscribe_op&) = delete;

	subscribe_op& operator=(subscribe_op&&) = default;
	subscribe_op& operator=(const subscribe_op&) = delete;

	using allocator_type = asio::associated_allocator_t<handler_type>;
	allocator_type get_allocator() const noexcept {
		return asio::get_associated_allocator(_handler);
	}

	using executor_type = asio::associated_executor_t<handler_type>;
	executor_type get_executor() const noexcept {
		return asio::get_associated_executor(_handler);
	}

	void perform(
		const std::vector<subscribe_topic>& topics,
		const subscribe_props& props
	) {
		_num_topics = topics.size();

		uint16_t packet_id = _svc_ptr->allocate_pid();
		if (packet_id == 0)
			return complete_immediate(client::error::pid_overrun, packet_id);

		if (_num_topics == 0)
			return complete_immediate(client::error::invalid_topic, packet_id);

		auto ec = validate_subscribe(topics, props);
		if (ec)
			return complete_immediate(ec, packet_id);

		auto subscribe = control_packet<allocator_type>::of(
			with_pid, get_allocator(),
			encoders::encode_subscribe, packet_id,
			topics, props
		);

		auto max_packet_size = _svc_ptr->connack_property(prop::maximum_packet_size)
				.value_or(default_max_send_size);
		if (subscribe.size() > max_packet_size)
			return complete_immediate(client::error::packet_too_large, packet_id);

		send_subscribe(std::move(subscribe));
	}

	void send_subscribe(control_packet<allocator_type> subscribe) {
		auto wire_data = subscribe.wire_data();
		_svc_ptr->async_send(
			wire_data,
			no_serial, send_flag::none,
			asio::prepend(
				std::move(*this), on_subscribe {}, std::move(subscribe)
			)
		);
	}

	void resend_subscribe(control_packet<allocator_type> subscribe) {
		if (_handler.cancelled() != asio::cancellation_type_t::none)
			return complete(
				asio::error::operation_aborted, subscribe.packet_id()
			);
		send_subscribe(std::move(subscribe));
	}

	void operator()(
		on_subscribe, control_packet<allocator_type> packet,
		error_code ec
	) {
		if (ec == asio::error::try_again)
			return resend_subscribe(std::move(packet));

		auto packet_id = packet.packet_id();

		if (ec)
			return complete(ec, packet_id);

		_svc_ptr->async_wait_reply(
			control_code_e::suback, packet_id,
			asio::prepend(std::move(*this), on_suback {}, std::move(packet))
		);
	}

	void operator()(
		on_suback, control_packet<allocator_type> packet,
		error_code ec, byte_citer first, byte_citer last
	) {
		if (ec == asio::error::try_again) // "resend unanswered"
			return resend_subscribe(std::move(packet));

		uint16_t packet_id = packet.packet_id();

		if (ec)
			return complete(ec, packet_id);

		auto suback = decoders::decode_suback(
			static_cast<uint32_t>(std::distance(first, last)), first
		);
		if (!suback.has_value()) {
			on_malformed_packet("Malformed SUBACK: cannot decode");
			return resend_subscribe(std::move(packet));
		}

		auto& [props, rcs] = *suback;
		auto reason_codes = to_reason_codes(std::move(rcs));
		if (reason_codes.size() != _num_topics) {
			on_malformed_packet(
				"Malformed SUBACK: does not contain a "
				"valid Reason Code for every Topic Filter"
			);
			return resend_subscribe(std::move(packet));
		}

		complete(
			ec, packet_id, std::move(reason_codes), std::move(props)
		);
	}

private:

	error_code validate_subscribe(
		const std::vector<subscribe_topic>& topics, const subscribe_props& props
	) const {
		error_code ec;
		for (const auto& topic: topics) {
			ec = validate_topic(topic);
			if (ec)
				return ec;
		}

		ec = validate_props(props);
		return ec;
	}

	error_code validate_topic(const subscribe_topic& topic) const {
		auto wildcard_available = _svc_ptr->connack_property(
			prop::wildcard_subscription_available
		).value_or(1);
		auto shared_available = _svc_ptr->connack_property(
			prop::shared_subscription_available
		).value_or(1);

		std::string_view topic_filter = topic.topic_filter;

		validation_result result = validation_result::valid;
		if (
			topic_filter.compare(0, shared_sub_prefix.size(), shared_sub_prefix) == 0
		) {
			if (!shared_available)
				return client::error::shared_subscription_not_available;

			result = validate_shared_topic_filter(topic_filter, wildcard_available);
		} else
			result = wildcard_available ?
				validate_topic_filter(topic_filter) :
				validate_topic_name(topic_filter);

		if (result == validation_result::invalid)
			return client::error::invalid_topic;
		if (!wildcard_available && result != validation_result::valid)
			return client::error::wildcard_subscription_not_available;
		return error_code {};
	}

	error_code validate_props(const subscribe_props& props) const {
		const auto& user_properties = props[prop::user_property];
		for (const auto& user_property: user_properties)
			if (!is_valid_string_pair(user_property))
				return client::error::malformed_packet;

		const auto& sub_id = props[prop::subscription_identifier];
		if (!sub_id.has_value())
			return error_code {};

		auto sub_id_available = _svc_ptr->connack_property(
			prop::subscription_identifier_available
		).value_or(1);

		if (!sub_id_available)
			return client::error::subscription_identifier_not_available;

		return (min_subscription_identifier <= *sub_id &&
			*sub_id <= max_subscription_identifier) ?
				error_code {} :
				client::error::malformed_packet;
	}

	static std::vector<reason_code> to_reason_codes(std::vector<uint8_t> codes) {
		std::vector<reason_code> ret;
		for (uint8_t code : codes) {
			auto rc = to_reason_code<reason_codes::category::suback>(code);
			if (rc)
				ret.push_back(*rc);
		}
		return ret;
	}

	void on_malformed_packet(const std::string& reason) {
		auto props = disconnect_props {};
		props[prop::reason_string] = reason;
		async_disconnect(
			disconnect_rc_e::malformed_packet, props, _svc_ptr,
			asio::detached
		);
	}

	void complete_immediate(error_code ec, uint16_t packet_id) {
		if (packet_id != 0)
			_svc_ptr->free_pid(packet_id);
		_handler.complete_immediate(
			ec, std::vector<reason_code>(_num_topics, reason_codes::empty),
			suback_props {}
		);
	}

	void complete(
		error_code ec, uint16_t packet_id,
		std::vector<reason_code> reason_codes = {}, suback_props props = {}
	) {
		if (reason_codes.empty() && _num_topics)
			reason_codes = std::vector<reason_code>(_num_topics, reason_codes::empty);

		if (!_svc_ptr->subscriptions_present()) {
			bool has_success_rc = std::any_of(
				reason_codes.cbegin(), reason_codes.cend(),
				[](const reason_code& rc) { return !rc; }
			);
			if (has_success_rc)
				_svc_ptr->subscriptions_present(true);
		}

		_svc_ptr->free_pid(packet_id);
		_handler.complete(ec, std::move(reason_codes), std::move(props));
	}
};

template <typename ClientService>
class initiate_async_subscribe {
	std::shared_ptr<ClientService> _svc_ptr;
public:
	explicit initiate_async_subscribe(std::shared_ptr<ClientService> svc_ptr) :
		_svc_ptr(std::move(svc_ptr))
	{}

	using executor_type = typename ClientService::executor_type;
	executor_type get_executor() const noexcept {
		return _svc_ptr->get_executor();
	}

	template <typename Handler>
	void operator()(
		Handler&& handler,
		const std::vector<subscribe_topic>& topics, const subscribe_props& props
	) {
		detail::subscribe_op { _svc_ptr, std::move(handler) }
			.perform(topics, props);
	}
};

} // end namespace async_mqtt5::detail

#endif // !ASYNC_MQTT5_SUBSCRIBE_OP_HPP
